{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lascar Example with CW Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Capturing is the same as usual. We'll capture a rather large amount of trace (3000) to show off the speed of lascar:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = \"OPENADC\"\n",
    "PLATFORM = \"CWLITEARM\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../Helper_Scripts/Setup.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "fw_path = \"../../hardware/victims/firmware/simpleserial-aes/simpleserial-aes-cwlitearm.hex\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# program the target\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Capture Traces\n",
    "from tqdm import tnrange\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "ktp = cw.ktp.Basic(target=target)\n",
    "\n",
    "traces = []\n",
    "textin = []\n",
    "keys = []\n",
    "N = 3000  # Number of traces\n",
    "project = cw.project(\"test.cwp\")\n",
    "for i in tnrange(N, desc='Capturing traces'):\n",
    "    # run aux stuff that should come before trace here\n",
    "\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "    textin.append(text)\n",
    "    keys.append(key)\n",
    "    \n",
    "    trace, resp = cw.capture_trace(scope, target, text, key)\n",
    "\n",
    "    traces.append(trace)\n",
    "\n",
    "#Convert traces to numpy arrays\n",
    "trace_array = np.asarray(traces)  # if you prefer to work with numpy array for number crunching\n",
    "textin_array = np.asarray(textin)\n",
    "known_keys = np.asarray(keys)  # for fixed key, these keys are all the same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import numpy as np\n",
    "\n",
    "project = cw.openProject(\"../stm32f415_lab.cwp\")\n",
    "\n",
    "start_point = 1312\n",
    "end_point = 1350\n",
    "\n",
    "tm = project.traceManager()\n",
    "\n",
    "trace_array = np.zeros( (tm.numTraces(), end_point - start_point))\n",
    "textin_array = np.zeros( (tm.numTraces(), len(tm.getTextin(0))), dtype=\"uint8\" )\n",
    "textout_array = np.zeros( (tm.numTraces(), len(tm.getTextout(0))), dtype=\"uint8\" )\n",
    "\n",
    "print (\"Copying %d traces of %d samples into memory\" % (tm.numTraces(), tm.numPoints()))\n",
    "for n in range(0, tm.numTraces()):\n",
    "    trace_array[n] = tm.getTrace(n)[start_point:end_point]\n",
    "    textin_array[n] = tm.getTextin(n)\n",
    "    textout_array[n] = tm.getTextout(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Lascar on Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing we will do is select a leakage function. The following is a few examples - the first is the standard sboxHW function, the next two are common for hardware crypto."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lascar import *\n",
    "from lascar.tools.aes import sbox, inv_sbox\n",
    "\n",
    "#The following leakage models copied from /chipwhisperer/analyzer/attacks/models/AES128_8bit.py and\n",
    "# massaged into Lascar Version\n",
    "\n",
    "def selection_function_sboxHW(byte):\n",
    "    # selection_with_guess function must take 2 arguments: value and guess\n",
    "    def selection_with_guess(value, guess):\n",
    "        return hamming(sbox[value[byte] ^ guess])\n",
    "    return selection_with_guess\n",
    "\n",
    "def selection_function_sboxInOutHD(byte):\n",
    "    # selection_with_guess function must take 2 arguments: value and guess\n",
    "    def selection_with_guess(value, guess):\n",
    "        return hamming(sbox[value[byte] ^ guess] ^ value[byte])\n",
    "    return selection_with_guess\n",
    "\n",
    "def selection_function_lastroundHD(byte):\n",
    "    # selection_with_guess function must take 2 arguments: value and guess\n",
    "    def selection_with_guess(value, guess):\n",
    "        INVSHIFT_undo = [0, 5, 10, 15, 4, 9, 14, 3, 8, 13, 2, 7, 12, 1, 6, 11]\n",
    "        st10 = value[INVSHIFT_undo[byte]]\n",
    "        st9 = inv_sbox[value[byte] ^ guess]\n",
    "        return hamming(st9 ^ st10)\n",
    "    return selection_with_guess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to make the engines. This requires us to select one of the above leakage functions and create a large array of them. Edit the following to change the leakage function you wish to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chipwhisperer.analyzer.attacks.models.AES128_8bit import LastroundStateDiff\n",
    "#Adjust this if needed - will ensure correct key/PGE highlighting is done!\n",
    "highlight_key = LastroundStateDiff().processKnownKey(project.traceManager().getKnownKey(0))\n",
    "#highlight_key = project.traceManager().getKnownKey(0)\n",
    "\n",
    "#Adjust this for actual attack used!\n",
    "guess_range = range(256)\n",
    "cpa_engines = [CpaEngine(\"cpa_%02d\" % i, selection_function_lastroundHD(i), guess_range) for i in range(16)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is alist of cpa engines now. See for example:\n",
    "#help(cpa_engines[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally run ONE of the following - the difference is the container either has the textin or textout. The correct setup depends on your leakage function assumptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Leakage models using plaintext (such as SBox output) require textin\n",
    "containter_textin =  TraceBatchContainer(trace_array, textin_array)\n",
    "session = Session(containter_textin, engines=cpa_engines).run(batch_size=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Leakage models using ciphertext (such as lastroundHD) require textout\n",
    "containter_textout =  TraceBatchContainer(trace_array, textout_array)\n",
    "session = Session(containter_textin, engines=cpa_engines).run(batch_size=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "import chipwhisperer as cw\n",
    "\n",
    "from chipwhisperer.analyzer.attacks._stats import DataTypeDiffs\n",
    "\n",
    "key = project.traceManager().getKnownKey(0)\n",
    "\n",
    "class LascarCWAttacks(object):\n",
    "    \n",
    "    def __init__(self, cpa_engines, highlight_key=None):  \n",
    "        dt = DataTypeDiffs()\n",
    "        for i in range(len(cpa_engines)):\n",
    "            results = cpa_engines[i].finalize()\n",
    "            dt.updateSubkey(i, results)\n",
    "\n",
    "        self.dt = dt\n",
    "        self.hlk = highlight_key\n",
    "\n",
    "    def display_pge(self):\n",
    "        cb = cw.getJupyterCallback(self)\n",
    "        cb()\n",
    "        \n",
    "    def getStatistics(self):\n",
    "        \"\"\"CW Interfae Function\"\"\"\n",
    "        return self.dt\n",
    "    \n",
    "    def knownKey(self):\n",
    "        \"\"\"CW Interface Function\"\"\"\n",
    "        if self.hlk is None: return [0]*16\n",
    "        \n",
    "        return self.hlk\n",
    "    \n",
    "    def getReportingInterval(self):\n",
    "        \"\"\"CW Interface Function\"\"\"\n",
    "        return 0\n",
    "    \n",
    "results = LascarCWAttacks(cpa_engines, highlight_key)\n",
    "results.display_pge()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unfinished Business"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "output_notebook()\n",
    "p = figure()\n",
    "for i in range(16):\n",
    "    results = cpa_engines[i].finalize()\n",
    "    xrange = range(len(results[0x2B]))\n",
    "    guess = abs(results).max(1).argmax()\n",
    "    print(\"Best Guess is {:02X} (Corr = {})\".format(guess, abs(results).max()))\n",
    "    p.line(xrange, results[guess])\n",
    "    \n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "\n",
    "output_notebook()\n",
    "xrange = range(len(results[0x2B]))\n",
    "print(xrange)\n",
    "print(len(results))\n",
    "p = figure()\n",
    "#p.line(xrange, traces[5],line_color='red')\n",
    "p.line(xrange, results[0x3C])\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mycontainer.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lascar.output.parse_results import apply_parse\n",
    "import pandas as pd\n",
    "from IPython.display import clear_output\n",
    "class JupyterOutputMethod(OutputMethod):\n",
    "    def __init__(self, *engines):\n",
    "        OutputMethod.__init__(self, engines)\n",
    "        self.parsed_subkeys = 0\n",
    "        self.results_list = [0] * 16\n",
    "    \n",
    "    def _update(self, engine, results):\n",
    "        engine.output_parser_mode = \"argmax\"\n",
    "        results_parsed = apply_parse(engine, abs(results))\n",
    "        display(self.parsed_subkeys)\n",
    "        if results_parsed is None:\n",
    "            return\n",
    "        idx = int(cpa_engine.name[-2:])\n",
    "        self.results_list[idx] = results_parsed\n",
    "        self.parsed_subkeys += 1\n",
    "        if self.parsed_subkeys > 15:\n",
    "            df = pd.DataFrame(self.results_list)\n",
    "            fd = df.transpose()\n",
    "            def formatter(stats):\n",
    "                return str(\"{}\".format(stats))\n",
    "            #clear_output(wait=True)\n",
    "            display(fd.head().style.format(formatter))\n",
    "            self.parsed_subkeys = 0\n",
    "        \n",
    "    def _finalize(self):\n",
    "        pass"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
